from dbnet.dbnet_infer import *

'''
    初始化全局参数
'''
pi = 3.14159265354
image_size = 1024

"""
    image's box_list index:
    0    1
    3    2
"""
def max_box_index(box_list,img):
    area_size = dict()
    count=0
    for i in box_list:
        x1,y1 = i[0]
        x2,y2 = i[2]
        # area_size[count] = sqrt(pow(abs(y2-y1),2) + pow(abs(x2-x1),2))
        area_size[count] = abs((y2-y1) * (x2-x1))
        count += 1
    area_size = sorted(area_size.items(),key= lambda d: d[1],reverse=True)
    return area_size[0][0]


'''
    先求出下标0-1之间的边长，可求出对应角度
    再与下标0-3边长对比，判断是否是长边，
    如果是短边，则+90度

    # TODO 缺点：可能是颠倒的文字，或是识别到其他文字
'''
def get_angle(img,points):
    angle = 0
    x_len = abs(points[0][0] - points[1][0])
    y_len = points[0][1] - points[1][1]
    third_len = sqrt(pow(x_len,2) + pow(abs(y_len),2))

    third_len2 = sqrt(pow(abs(points[0][0] - points[3][0]),2) 
                    + pow(abs(points[0][1] - points[3][1]),2))
    if(y_len < 0):
        angle -= asin( abs(y_len) / third_len)*180/pi
    else:
        angle += asin( abs(y_len) / third_len)*180/pi
    
    print("角度:",angle)
    if(third_len < third_len2):
        angle += 90      
    # cv2.putText(img,"Angle:"+str(angle),(30,50),
    #                     cv2.FONT_HERSHEY_SIMPLEX, 2, (0,0,255), 3)

    return angle

def adjust_angle(img,angle):
    rotated = img
    if(abs(angle) > 5):
        h, w, _ = img.shape
        center = (w // 2, h // 2)
        M = cv2.getRotationMatrix2D(center, -angle, 1.0)
        rotated = cv2.warpAffine(img, M, (w, h), flags=cv2.INTER_CUBIC, borderMode=cv2.BORDER_REPLICATE)
    return rotated


'''
    单次处理测试用例
'''
def angle_correct_test(path,name):
    img = cv2.imread(path)
    img= cv2.resize(img,(image_size,image_size))
    cv2.imwrite("/home/meroke/code/detection/result/{}origin.jpg".format(name), img)
    # 返回识别框的坐标点列表
    box_list, score_list = text_handle.process(img,image_size)
    if(box_list is None):
        print("Error with {}".format(path))
        return
    # 获取最大识别框下标
    max_index = max_box_index(box_list,img)

    # 绘制展示最大识别框
    img = draw_max_bbox(img, box_list[max_index],(0,0,255))

    # 根据固定下标0 1 两点 计算图片应该旋转的角度
    angle = get_angle(img,box_list[max_index])
    cv2.imwrite("/home/meroke/code/detection/result/{}_del_before.jpg".format(name), img)

    # 图像矫正
    img = adjust_angle(img,angle)

    h, w, _ = img.shape
    center = (w // 2, h // 2)
    M = cv2.getRotationMatrix2D(center, -180, 1.0)
    rotated = cv2.warpAffine(img, M, (w, h), flags=cv2.INTER_CUBIC, borderMode=cv2.BORDER_REPLICATE)
    cv2.imwrite("/home/meroke/code/detection/result/{}_del_after.jpg".format(name), rotated)
    name = name.split('.')[0]
    cv2.imwrite("/home/meroke/code/detection/result/{}_del.jpg".format(name), img)


# 识别模型载入
<<<<<<< HEAD
text_handle = DBNET(MODEL_PATH="./models/dbnet.onnx")
=======
pwd = os.popen("pwd")
path = pwd.read().strip()
if path > "/home/eaibot/2022_jsjds/detection_2022":
    path = path[:-4]
model_path = path + "/ocr/models/dbnet.onnx"
text_handle = DBNET(MODEL_PATH=model_path)
>>>>>>> origin/develop
'''
    批量处理测试用例
'''
if __name__ == "__main__":
    floder_path = "/home/meroke/code/detection/result/error/"
    for i in os.listdir(floder_path):
        if "jpg" in i:
            if "del" in i:
                continue
            angle_correct_test(floder_path + i,i)
            print(i + " done!")